package com.yifeng.repo.toolkit.mq;

import com.alibaba.cloud.stream.binder.rocketmq.constant.RocketMQConst;
import com.google.common.base.Strings;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.client.producer.LocalTransactionState;
import org.apache.rocketmq.client.producer.TransactionListener;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageExt;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * Created by daibing on 2022/9/22.
 */
@Slf4j
public class TempTransactionListener implements TransactionListener {
    /**
     * 注意：这2个变量仅用做测试使用，上线生产环境必须去掉字段和相应的测试逻辑
     */
    private final AtomicInteger execCount = new AtomicInteger(0);
    private final AtomicInteger checkCount = new AtomicInteger(0);

    /**
     * Execute local transaction.
     * @param msg messages
     * @param arg message args
     * @return Transaction state
     */
    @Override
    public LocalTransactionState executeLocalTransaction(Message msg, Object arg) {
        // Check header
        String userTransactionalArgs = msg.getProperty(RocketMQConst.USER_TRANSACTIONAL_ARGS);
        String transactionId = msg.getProperty(RocketMQConst.Headers.TRANSACTION_ID);
        String key = msg.getProperty("KEYS");
        if (Strings.isNullOrEmpty(userTransactionalArgs) || Strings.isNullOrEmpty(transactionId) || Strings.isNullOrEmpty(key)) {
            log.info("tx listener execute status: state={}, userTransactionalArgs={}, transactionId={}, key={}",
                    LocalTransactionState.ROLLBACK_MESSAGE, userTransactionalArgs, transactionId, key);
            return LocalTransactionState.ROLLBACK_MESSAGE;
        }

        // order return tx state
        LocalTransactionState state = LocalTransactionState.values()[execCount.getAndIncrement() % LocalTransactionState.values().length];
        log.info("tx listener execute status: state={}, userTransactionalArgs={}, transactionId={}, key={}",
                state, userTransactionalArgs, transactionId, key);
        return state;
    }

    /**
     * MQ check back local transaction states.
     * @param msg messages
     * @return Transaction state
     */
    @Override
    public LocalTransactionState checkLocalTransaction(MessageExt msg) {
        // Check header
        String userTransactionalArgs = msg.getProperty(RocketMQConst.USER_TRANSACTIONAL_ARGS);
        String transactionId = msg.getProperty(RocketMQConst.Headers.TRANSACTION_ID);
        String key = msg.getProperty("KEYS");
        if (Strings.isNullOrEmpty(userTransactionalArgs) || Strings.isNullOrEmpty(transactionId) || Strings.isNullOrEmpty(key)) {
            log.info("tx listener check status: state={}, userTransactionalArgs={}, transactionId={}, key={}",
                    LocalTransactionState.ROLLBACK_MESSAGE, userTransactionalArgs, transactionId, key);
            return LocalTransactionState.ROLLBACK_MESSAGE;
        }

        // order return tx state
        LocalTransactionState state = checkCount.getAndIncrement() % 2 == 0 ? LocalTransactionState.COMMIT_MESSAGE : LocalTransactionState.ROLLBACK_MESSAGE;
        log.info("tx listener check status: state={}, userTransactionalArgs={}, transactionId={}, key={}",
                state, userTransactionalArgs, transactionId, key);
        return state;
    }
}
